import logging
import httpx
from enum import Enum
import asyncio
from pydantic import BaseModel, Field
import typing as T
from .Global import ENV


class GenoratorImage:
    BASE_URL = "https://dashscope.aliyuncs.com/api/v1"
    TEXT2IMAGE_URL = f"{BASE_URL}/services/aigc/text2image/image-synthesis"

    AUTH_HEADERS = {
        "Authorization": f"Bearer {ENV.API_KEY}",
    }
    TEXT2IMAGE_HEADERS = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {ENV.API_KEY}",
        "Content-Type": "application/json",
    }
    client = httpx.AsyncClient()

    class RESULT_TYPE(str, Enum):
        SUCCESS = "succeeded"
        FAILED = "failed"
        TIMEOUT = "timeout"
        CREATE_ERROR = "create_error"

    class RESULT(BaseModel):
        type: T.Literal["succeeded", "failed", "timeout", "create_error"] = Field(
            description="任务状态"
        )
        results: T.Optional[T.List[T.Dict[str, str]]] = Field(
            default=None, description="生成图片的url"
        )
        code: T.Optional[str] = Field(default=None, description="错误码")
        message: T.Optional[str] = Field(default=None, description="错误信息")

    @classmethod
    async def create_async_task(cls, params: dict):

        try:
            # 使用任务ID获取结果
            response = await cls.client.post(
                cls.TEXT2IMAGE_URL, headers=cls.TEXT2IMAGE_HEADERS, json=params
            )
            response.raise_for_status()
            body = response.json()
            if not body.get("output"):
                return cls.RESULT_TYPE.CREATE_ERROR, body
            return cls.RESULT_TYPE.SUCCESS, body
        except httpx.HTTPStatusError as e:
            logging.info(f"发生HTTP错误: {e}")
            raise Exception(f"发生HTTP错误: {e}")
        except httpx.RequestError as e:
            logging.info(f"发生请求错误: {e}")
            raise Exception(f"发生请求错误: {e}")

    @classmethod
    async def get_task_result(cls, task_id: str):
        try:
            response = await cls.client.get(
                f"{cls.BASE_URL}/tasks/{task_id}", headers=cls.AUTH_HEADERS
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as e:
            logging.info(f"发生HTTP错误: {e}")
            raise Exception(f"发生HTTP错误: {e}")

    @classmethod
    async def wait_async_task(
        cls, task_id: str, timeout: float = 60.0, interval: float = 2.5
    ):
        start_time = asyncio.get_running_loop().time()

        while True:
            result = await cls.get_task_result(task_id)
            output = result.get("output")
            if output.get("task_status").lower() == cls.RESULT_TYPE.SUCCESS:
                return cls.RESULT_TYPE.SUCCESS, output.get("results")
            elif output.get("task_status").lower() == cls.RESULT_TYPE.FAILED:
                return cls.RESULT_TYPE.FAILED, output
            elif asyncio.get_running_loop().time() - start_time > timeout:
                return cls.RESULT_TYPE.TIMEOUT, None
            else:
                await asyncio.sleep(interval)

    @classmethod
    async def create_task_and_wait(
        cls, params: dict, timeout: float = 60.0, interval: float = 2.5
    ) -> RESULT:
        tag, result = await cls.create_async_task(params)
        if tag == cls.RESULT_TYPE.CREATE_ERROR:
            return cls.RESULT(
                type="create_error",
                message=result.get("message"),
                code=result.get("code"),
            )
        task_id = result.get("output").get("task_id")
        tag, results = await cls.wait_async_task(task_id, timeout, interval)
        if tag == cls.RESULT_TYPE.FAILED:
            return cls.RESULT(
                type="failed", message=results.get("message"), code=results.get("code")
            )
        elif tag == cls.RESULT_TYPE.TIMEOUT:
            return cls.RESULT(type="timeout")
        else:
            return cls.RESULT(type="succeeded", results=results)
